package confirms;

import com.rabbitmq.client.Channel;
import com.rabbitmq.client.ConfirmListener;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;
import constant.Constants;

import java.io.IOException;
import java.util.Collections;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.TimeoutException;

public class Producer {
    public static int MESSAGE_COUNT = 200;
    public static int WAIT_TIME = 5000;
    public static int BATCH_SIZE = 50;

    // 建立连接
    public static Connection createConnection() throws IOException, TimeoutException {
        // 1. 创建连接工厂
        ConnectionFactory factory = new ConnectionFactory();
        // 2. 设置参数
        factory.setHost(Constants.HOST); // ip 的默认值为 localhost
        factory.setPort(Constants.PORT); // 默认值为 5672
        factory.setVirtualHost(Constants.VIRTUAL_HOST); // 虚拟主机，默认值为 /
        // 账号
        factory.setUsername(Constants.USER_NAME); // 用户名，默认为 guest
        factory.setPassword(Constants.USER_PASSWORD); // 密码，默认为 guest
        // 3. 创建连接 Connection
        Connection connection = factory.newConnection(); // 需要处理异常，在此处直接抛出，并不进行处理
        return connection;
    }

    // 单独确认模式
    public static void publishMessageIndividually() {
        try (Connection connection = createConnection()){
            // 开启信道
            Channel channel = connection.createChannel();
            // 开启信道确认模式
            channel.confirmSelect();
            // 声明队列
            channel.queueDeclare(Constants.PUBLISH_CONFIRMS_QUEUE_1, true, false, false, null);
            // 记录开始时间
            long startTime = System.currentTimeMillis();
            // 发送消息
            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String message = "message " + i;
                channel.basicPublish("", Constants.PUBLISH_CONFIRMS_QUEUE_1, null, message.getBytes());
                // 等待确认
                channel.waitForConfirmsOrDie(WAIT_TIME);
            }
            // 记录结束时间
            long endTime = System.currentTimeMillis();
            System.out.printf("publish %d messages individually in %d ms\n", MESSAGE_COUNT, endTime - startTime);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } catch (TimeoutException e) {
            throw new RuntimeException(e);
        }
    }

    // 批量确认模式
    public static void publishMessageInBatches() {
        try (Connection connection = createConnection()){
            // 开启信道
            Channel channel = connection.createChannel();
            // 设置为 confirm 模式
            channel.confirmSelect();
            // 声明队列
            channel.queueDeclare(Constants.PUBLISH_CONFIRMS_QUEUE_2, true, false, false, null);
            // 发送消息
            int messageCount = 0;
            long startTime = System.currentTimeMillis();
            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String message = "message " + i;
                channel.basicPublish("", Constants.PUBLISH_CONFIRMS_QUEUE_2, null, message.getBytes());
                messageCount++;
                // 批量确认
                if(messageCount == BATCH_SIZE) {
                    channel.waitForConfirms(WAIT_TIME);
                    messageCount = 0;
                }
            }
            // 消息发送完，若还有未确认消息，则进行最后的确认
            if (messageCount > 0) {
                channel.waitForConfirms(WAIT_TIME);
            }
            long endTime = System.currentTimeMillis();
            System.out.printf("publish %d messages in batch in %d ms\n", MESSAGE_COUNT, endTime - startTime);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } catch (TimeoutException e) {
            throw new RuntimeException(e);
        }
    }

    // 异步确认模式
    public static void handlePublishConfirmAsynchronously() {
        try (Connection connection = createConnection()){
            // 开启信道
            Channel channel = connection.createChannel();
            // 设置 confirm 模式
            channel.confirmSelect();
            // 声明队列
            channel.queueDeclare(Constants.PUBLISH_CONFIRMS_QUEUE_3, false, false, true, null);
            // 使用有序集合来存储未确认的消息
            SortedSet<Long> confirmSet = Collections.synchronizedSortedSet(new TreeSet<>());
            // 进行确认
            channel.addConfirmListener(new ConfirmListener() {
                @Override
                public void handleAck(long deliveryTag, boolean multiple) throws IOException {
                    // 若为批量确认，则删除确认序号前所有元素
                    if (multiple) {
                        confirmSet.headSet(deliveryTag + 1).clear();
                    } else {
                        confirmSet.remove(deliveryTag);
                    }
                }

                @Override
                public void handleNack(long deliveryTag, boolean multiple) throws IOException {
                    // 处理失败，消息重发...
                    // 若为批量确认，则删除确认序号前所有元素
                    if (multiple) {
                        confirmSet.headSet(deliveryTag + 1).clear();
                    } else {
                        confirmSet.remove(deliveryTag);
                    }
                }
            });
            // 记录开始时间
            long startTime = System.currentTimeMillis();
            // 发送消息
            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String message = "message " + i;
                // 获取序号
                long nextPublishSeqNo = channel.getNextPublishSeqNo();
                channel.basicPublish("", Constants.PUBLISH_CONFIRMS_QUEUE_3, null, message.getBytes());
                // 存储序号
                confirmSet.add(nextPublishSeqNo);
            }
            // 消息确认完毕
            while (!confirmSet.isEmpty()) {
                Thread.sleep(10);
            }
            // 记录结束时间
            long endTime = System.currentTimeMillis();
            System.out.printf("publish %d messages and handled confirms asynchronously in %d ms\n", MESSAGE_COUNT, endTime - startTime);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } catch (TimeoutException e) {
            throw new RuntimeException(e);
        }
    }

    public static void main(String[] args) {
        // 单独确认模式
        publishMessageIndividually();
        // 批量确认模式
         publishMessageInBatches();
        // 异步确认
         handlePublishConfirmAsynchronously();
    }
}
